import tensorflow.compat.v1 as tf
tf.compat.v1.disable_eager_execution()
import cv2 as cv
import numpy as np

if __name__ == "__main__":
    test_pb_model = True
    test_tflite_model = False
    read_cahnge_graph = False

    pb_model_path = "data\model.pb"
    input_node_name = "images"
    output_node_name = "all_class_predictions_with_background"

    src_img = cv.imread("img\IMG_kitti_0011_000187.png")
    cv.imwrite("src.jpg", src_img)

    src_img = cv.resize(src_img, (960, 540))
    src_img = cv.cvtColor(src_img, cv.COLOR_BGR2YCrCb)[:, :, 0]
    src_img = src_img / 127.5 - 1
    src_img = src_img.astype("float32")

    src_img = src_img.reshape((1, 540, 960, 1))

    if test_tflite_model:
        interpreter = tf.lite.Interpreter(tflite_model_path)
        interpreter.allocate_tensors()

        input_details = interpreter.get_input_details()
        # print(str(input_details))
        output_details = interpreter.get_output_details()

        interpreter.set_tensor(input_details[0]["index"], src_img)

        interpreter.invoke()
        output_data = interpreter.get_tensor(output_details[0]["index"])

        result = output_data[0]

        result = (result + 1) * 127.5
        result[result > 255] = 255
        result[result < 0] = 0
        result = result.astype(np.uint8)
        cv.imshow("result", result)
        cv.imwrite("result.jpg", result)
        cv.waitKey()
    if test_pb_model:
        src_img = cv.imread("img\IMG_kitti_0011_000187.png")
        src_img = cv.resize(src_img, (960, 540))
        src_img = cv.cvtColor(src_img, cv.COLOR_BGR2YCrCb)[:, :, 0]
        src_img = src_img / 127.5 - 1
        src_img = src_img.astype("float32")
        src_img = src_img.reshape((1, 540, 960, 1))
        input_image = tf.placeholder(tf.float32, (1, 540, 960, 1))

        with open(pb_model_path, "rb") as f:
            graph_def = tf.GraphDef()
            graph_def.ParseFromString(f.read())
            out_result = tf.import_graph_def(
                graph_def,
                input_map={"images:0": input_image},
                return_elements=["map/TensorArrayUnstack/range:0"])
        sess = tf.Session()
        result = sess.run(out_result, feed_dict={input_image: src_img})

        result = result[0][0]
        result = (result + 1) * 127.5
        # result[result > 255] = 255
        # result[result < 0] = 0
        result = result.astype(np.uint8)

        cv.imshow("result", result)
        cv.waitKey(0)
    if read_cahnge_graph:
        gf = tf.GraphDef()
        gf.ParseFromString(open(pb_model_path, "rb").read())
        for n in gf.node:
            print(n.name + " ===> " + n.op)
